package com.xp.pgoj.controller;

import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import com.xp.pgoj.common.BaseResponse;
import com.xp.pgoj.common.ResultUtils;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;

import static com.xp.pgoj.constant.AIUrlConstant.ASYNC_URL;
import static com.xp.pgoj.constant.AIUrlConstant.HEADERS;

/**
 * @author <a href="https://github.com/Fzs0">逢左使</a>
 * @create 2024/10/01 03:02
 * @Description: AI录题接口
 */
@RestController
@RequestMapping("/ai")
@Slf4j
public class AIController {

    @Value("${glm.apiSecretKey}")
    private String apiSecretKey;

    private final ExecutorService executorService = Executors.newSingleThreadExecutor();

    @PostMapping("/question")
    public BaseResponse<String> getAnswerByAI(String content) {
        ClientV4 client = new ClientV4.Builder(apiSecretKey).build();
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), content);
        messages.add(chatMessage);

        long l = System.currentTimeMillis();
        String requestId = String.valueOf(l);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethodAsync)
                .requestId(requestId)
                .messages(messages)
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        System.out.println("modelAPIResponse output ID  ====>" + invokeModelApiResp.getData().getId());
        System.out.println("modelAPIResponse output  ====>" + invokeModelApiResp.getData());

        // 根据异步url来获取返回值
        String getResponseByAsyncUrl = ASYNC_URL + String.valueOf(invokeModelApiResp.getData().getId());
        AtomicReference<String> responseContent = new AtomicReference<>("");

        CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
            while (true) {
                String response = HttpUtil.get(getResponseByAsyncUrl, HEADERS);
                if (handleResponse(response, responseContent)) {
                    break;
                }
                try {
                    Thread.sleep(5000); // 每5秒检查一次
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        }, executorService);

        try {
            future.get(60, TimeUnit.SECONDS); // 等待最多30秒
        } catch (InterruptedException | ExecutionException | TimeoutException e) {
            e.printStackTrace();
        }

        return ResultUtils.success(responseContent.get());
    }

    private boolean handleResponse(String jsonResponse, AtomicReference<String> responseContent) {
        JSONObject jsonObject = new JSONObject(jsonResponse);
        String taskStatus = jsonObject.getStr("task_status");

        if ("SUCCESS".equals(taskStatus)) {
            String content = jsonObject.getJSONArray("choices").getJSONObject(0).getJSONObject("message").getStr("content");
            System.out.println("任务成功，内容是: " + content);
            responseContent.set(content);
            return true;
        } else if ("FAILED".equals(taskStatus)) {
            System.out.println("任务失败，停止监控。");
            responseContent.set("任务失败");
            return true;
        } else {
            System.out.println("任务状态仍在处理中: " + taskStatus);
            return false;
        }
    }
}